import torch
import torch.nn as nn


class CNN(nn.Module):
    def __init__(self, in_channel):
        super(CNN, self).__init__()
        self.block1 = nn.Sequential(
            nn.Conv2d(in_channels=in_channel, out_channels=32, kernel_size=3),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(2)
        )
        self.block2 = nn.Sequential(
            nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(2)
        )
        self.block3 = nn.Sequential(
            nn.Conv2d(in_channels=64, out_channels=64, kernel_size=3),
            nn.ReLU(inplace=True),
        )
        self.block4 = nn.Sequential(
            nn.Linear(in_features=1544192, out_features=64),
            nn.ReLU(inplace=True),
        )
        self.fc = nn.Linear(in_features=64, out_features=10)
        self.flatten = nn.Flatten()

    def forward(self, x):
        f = self.block1(x)

        f1 = self.block2(f)
        f2 = self.block3(f1)

        f3 = self.flatten(f2)
        # print(f3.shape)
        f4 = self.block4(f3)
        out = self.fc(f4)
        return out


if __name__ == '__main__':
    device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
    data = torch.rand(1, 4, 480, 848)
    data = data.to(device)
    model = CNN(in_channel=4).to(device)
    out = model(data)
    print(out.shape)
